"""
联邦学习多源数据融合策略
策略特点：
1. 使用联邦学习框架聚合多数据源模型
2. 保护数据隐私的同时提升预测精度
3. 支持跨市场、跨资产类别的联合建模
"""

import torch
import numpy as np
from torch import nn
from collections import OrderedDict

class LocalModel(nn.Module):
    """本地数据源模型"""
    def __init__(self, input_dim=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, 1)
        )
        
    def forward(self, x):
        return self.net(x)

class FederatedStrategy:
    def __init__(self, data_sources):
        """
        data_sources: 各数据源名称列表 
        """
        self.data_sources = data_sources
        self.local_models = {name: LocalModel() for name in data_sources}
        self.global_model = LocalModel()
        
    def aggregate_models(self):
        """联邦平均聚合各本地模型"""
        global_dict = self.global_model.state_dict()
        
        # 初始化累加器
        for k in global_dict.keys():
            global_dict[k] = torch.zeros_like(global_dict[k])
            
        # 加权平均
        total_samples = sum([m['samples'] for m in self.local_models.values()])
        for name, model_info in self.local_models.items():
            local_dict = model_info['model'].state_dict()
            weight = model_info['samples'] / total_samples
            for key in global_dict.keys():
                global_dict[key] += local_dict[key] * weight
                
        # 更新全局模型
        self.global_model.load_state_dict(global_dict)
        
    def train_local(self, data_source, dataset, epochs=10):
        """在指定数据源上训练本地模型"""
        model = self.local_models[data_source]['model']
        optimizer = torch.optim.Adam(model.parameters())
        criterion = nn.MSELoss()
        
        for _ in range(epochs):
            for x, y in dataset:
                optimizer.zero_grad()
                pred = model(x)
                loss = criterion(pred, y)
                loss.backward()
                optimizer.step()
                
        # 记录样本数
        self.local_models[data_source]['samples'] = len(dataset)
        
    def predict(self, x, source='global'):
        """使用指定模型进行预测"""
        model = self.global_model if source == 'global' else self.local_models[source]['model']
        with torch.no_grad():
            return model(x).item()

if __name__ == '__main__':
    # 示例用法
    sources = ['broker1', 'broker2', 'institutional']
    strategy = FederatedStrategy(sources)
    
    # 模拟训练数据
    torch.manual_seed(42)
    datasets = {
        'broker1': [(torch.randn(10), torch.randn(1)) for _ in range(100)],
        'broker2': [(torch.randn(10), torch.randn(1)) for _ in range(150)], 
        'institutional': [(torch.randn(10), torch.randn(1)) for _ in range(200)]
    }
    
    # 本地训练
    for source in sources:
        strategy.train_local(source, datasets[source])
    
    # 联邦聚合
    strategy.aggregate_models()
    
    # 测试预测
    test_input = torch.randn(10)
    print("全局模型预测:", strategy.predict(test_input))
    print("broker1模型预测:", strategy.predict(test_input, 'broker1'))